# 导入所需的库
import torch
from torch import nn
import torch.optim as optim

# 1. 准备数据
# 假设我们有以下模拟数据，实际项目中应从文件或数据库读取
# X是特征，Y是目标变量
X = torch.tensor([[1.0], [2.0], [3.0], [4.0]], requires_grad=True)
Y = torch.tensor([[2.0], [4.0], [6.0], [8.0]], dtype=torch.float)

# 将数据分为训练集和验证集（这里为了简化不进行划分）

# 2. 设计模型 - 线性回归模型只有一个线性层（没有激活函数）
class LinearRegression(nn.Module):
    def __init__(self, input_dim=1):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(input_dim, 1)  # 输出维度为1，对应一个目标值

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = LinearRegression()

# 3. 损失函数 - 我们通常会选择均方误差作为线性回归的损失函数
criterion = nn.MSELoss()  # 均方误差损失函数

# 4. 优化器 - 使用随机梯度下降等优化算法更新模型参数
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 学习率为0.01

# 5. 训练循环
num_epochs = 100  # 总迭代次数
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(X)
    loss = criterion(outputs, Y)  # 计算损失

    # 反向传播和优化
    optimizer.zero_grad()  # 清零梯度
    loss.backward()  # 反向传播计算梯度
    optimizer.step()  # 根据梯度更新参数

    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 6. 测试模型
with torch.no_grad():
    predictions = model(X)
    print('Predictions:', predictions)
